"""

从demo_reg.py保存的模型结果文件，model_reg.pkl
加载模型后，直接用于推理！

"""
import torch
from load_data import load_data, split_train_test

X, Y = load_data("./housing.data")
X_train, X_test, Y_train, Y_test = split_train_test(X, Y)
print("X_train.shape: ", X_train.shape)
print("X_test.shape: ", X_test.shape)
print("Y_train.shape: ", Y_train.shape)
print("Y_test.shape: ", Y_test.shape)


# net
class Net(torch.nn.Module):
    def __init__(self, n_feature, n_out):
        super().__init__()

        # 不使用隐藏层，直接建立回归
        # self.predict = torch.nn.Linear(n_feature, n_out)

        # 创建一个隐藏层
        self.hidden = torch.nn.Linear(n_feature, 100)
        self.predict = torch.nn.Linear(100, n_out)

    def forward(self, x):
        out = self.hidden(x)
        out = torch.relu(out)

        out = self.predict(out)

        return out


net = torch.load("model/model_reg.pkl")
# 损失函数
loss_func = torch.nn.MSELoss()

# 推理：
x_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(Y_test, dtype=torch.float32)
pred = net.forward(x_test)
pred = torch.squeeze(pred)
loss_test = loss_func(pred, y_test) * 0.001
print("loss_test: {}".format(loss_test))
